# Copyright (c) 2023-2024, ETH Zurich (Robotics Systems Lab)
# Author: Pascal Roth
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

# python
import os
from dataclasses import dataclass, field
from typing import List, Optional

# orbit-assets
from omni.isaac.assets import ASSETS_RESOURCES_DIR


@dataclass
class ANYmalEvaluatorConfig:
    # termination conditions
    max_time = 60  # seconds
    max_remain_time = 5  # seconds

    # cost map file
    cost_map_dir: Optional[str] = None  # "/home/pascal/viplanner/imperative_learning/data/cost_maps"
    cost_map_name: Optional[str] = None  # "2n8kARJN3HM_cost_map_long_2"

    # use previous results or previous generated waypoints
    use_prev_results: bool = True
    use_existing_explored_waypoints: bool = True
    handcrafted_waypoint_file: Optional[str] = None  # "2n8kARJN3HM_waypoints_long"  # _2
    waypoint_dir: Optional[str] = "/home/pascal/viplanner/imperative_learning/data/waypoints"
    # NOTE: can either load waypoint generated by the waypoint extension (define file under handcrafted_waypoint_file)
    #       or load previously explored waypoints (define file under handcrafted_waypoint_file), if neither env will
    #       be explored
    repeat_waypoints: Optional[int] = 50  # number of times to repeat the waypoints, create distribution of results

    # waypoint exploration parameters
    num_pairs: int = 500  # number of start-goal pairs to explore
    min_goal_dist: int = 5  # meters
    max_goal_dist: int = 15  # meters
    num_connections: int = 3  # number of connections when building the graph of all samples
    seed: int = 1  # random seed for collection of start-goal points

    # multi model
    multi_model: bool = False
    models: List[str] = field(default_factory=list)
    save_dir: str = os.path.join(ASSETS_RESOURCES_DIR, "vip_models")

    # intermediate result saving period
    save_period: int = 10  # waypoints

    @property
    def waypoint_file(self):
        assert all([self.waypoint_dir, self.handcrafted_waypoint_file]), "Waypoint file not specified"
        return os.path.join(self.waypoint_dir, f"{self.handcrafted_waypoint_file}.json")


# EoF
